# /home/iris/Project/python/bio/data/survial.csv
import pandas as pd
import matplotlib.pyplot as plt
from lifelines import KaplanMeierFitter
from lifelines.statistics import logrank_test
# df.shape[1] 返回列数
# df.shape[0] 返回行数

csv_file = "./survial_COHL.csv"
gene = 'AHR'
csv_data = pd.read_csv(csv_file, low_memory = False,header=0, usecols=['sample', 'OS',"OS.time",gene])
# csv_data = csv_data.sort(columns = [gene],axis = 0,ascending = True)
csv_data = csv_data.sort_values(by=gene,ascending=True)
# 取中位数
l = csv_data.shape[0]
if (l % 2) == 0:
    m = int(l/2)
else:
    m = int((l-1)/2)

group = []
for i in range(0,m):
    group.append('low')
for i in range(m,l):
    group.append("high")
csv_data['group'] = group
# csv_data['group'][m:l] = "high"
# csv_data['group'][0:m] = "low"


km = KaplanMeierFitter()
T = csv_data['OS.time']
E = csv_data['OS']
high = (csv_data['group'] == 'high')
ax = plt.subplot(111)

km.fit(T[high], event_observed=E[high], label="High level")
km.plot(ax=ax)
km.fit(T[~high], event_observed=E[~high], label="Low level")
km.plot(ax=ax)
plt.show()

# lr = logrank_test(T[high], T[~high], E[high], E[~high], alpha=.99)
# print('a:'+str(lr.p_value))
